/*
 * AqueousPhase.cpp
 *
 *  Created on: 23 May 2011
 *      Author: allan
 */

#include "AqueousPhase.h"

// GeoReact includes
#include "ActivityCoefficient.h"
#include "Utils.h"

// C++ includes
#include <algorithm>
#include <numeric>

AqueousPhase::AqueousPhase() 
{}

AqueousPhase::~AqueousPhase() 
{}

void AqueousPhase::AddSpecies(string nameOfSpecies)
{
	Phase::AddSpecies(nameOfSpecies);
	
	// store the electrical charge of the ion
	z.push_back(ElectricalCharge(nameOfSpecies));
	
	// store the default value of the Setschenow constant for the species, which is used only if it is neutral
	setschenowConstants.push_back(0.1);
	
	// check if the species is an ion
	if(z.back() != 0)
	{
		ions.push_back(GetNumOfSpecies() - 1);
		
		// set the HKF aqueous activity model for the ionic species
		activityModels.push_back(AQUEOUSACTIVITYMODEL_HKF_ION);
		
		// set the effective electrostatic radius of the ion
		effRadii.push_back(EffectiveElectrostaticRadiusOfIon(nameOfSpecies));
	}
	// else, check if the species is the water species
	else if(nameOfSpecies == "H2O(l)")
	{
		indexOfWater = GetNumOfSpecies() - 1;
		
		// set the HKF aqueous activity model for the water species
		activityModels.push_back(AQUEOUSACTIVITYMODEL_HKF_WATER);
	}
	// if the above comparisons failed, than the species is neutral but not water
	else
	{
		// set the Setschenow aqueous activity model for the neutral species
		activityModels.push_back(AQUEOUSACTIVITYMODEL_SETSCHENOW);
	}
	
	// setting the mapping species to ions, i.e., given the index of the species, you can retrieve the index of the ion (if not ion, unsigned(-1) is retrieved)
	if(z.back() != 0)
		species2ions.push_back(GetNumOfIons() - 1);
	else
		species2ions.push_back(unsigned(-1));
}

void AqueousPhase::SetElectrolyte(string electrolyte, const vector<pair<string,double>>& electrolyteIons)
{
	electrolytes.push_back(GetIndexOfSpecies(electrolyte));
	
	vector<pair<unsigned,double>> electrolyteIonsAux;
	
	for(unsigned iter = 0; iter < electrolyteIons.size(); ++iter)
	{
		const string nameOfIon = electrolyteIons[iter].first;
		
		const unsigned iIon = GetIndexOfIon(nameOfIon);
		
		const double stoichiometryOfIon = electrolyteIons[iter].second;
		
		electrolyteIonsAux.push_back({iIon, stoichiometryOfIon});
	}
	
	ionsOfElectrolytes.push_back(electrolyteIonsAux);
}

void AqueousPhase::SetActivityCoefficientModelHKF()
{
	for(unsigned iIon = 0; iIon < GetNumOfIons(); ++iIon)
		activityModels[GetIndexOfIon(iIon)] = AQUEOUSACTIVITYMODEL_HKF_ION;
	
	activityModels[indexOfWater] = AQUEOUSACTIVITYMODEL_HKF_WATER;
}

void AqueousPhase::SetActivityCoefficientModelSetschenow(string nameOfSpecies, double setschenow)
{
	unsigned iSpecies = GetIndexOfSpecies(nameOfSpecies);
	
	setschenowConstants[iSpecies] = setschenow;
	
	activityModels[iSpecies] = AQUEOUSACTIVITYMODEL_SETSCHENOW;
}

const unsigned AqueousPhase::GetNumOfIons() const
{
	return ions.size();
}

const unsigned AqueousPhase::GetNumOfElectrolytes() const
{
	return electrolytes.size();
}

const unsigned AqueousPhase::GetIndexOfWater() const
{
	return indexOfWater;
}

const unsigned AqueousPhase::GetIndexOfIon(unsigned iIon) const
{
	return ions[iIon];
}

const unsigned AqueousPhase::GetIndexOfIon(string nameOfIon) const
{
	for(unsigned iIon = 0; iIon < GetNumOfIons(); ++iIon)
		if(GetNameOfSpecies(GetIndexOfIon(iIon)) == nameOfIon)
			return iIon;
	
	return unsigned(-1);
}

const unsigned AqueousPhase::GetIndexOfElectrolyte(unsigned iElectrolyte) const
{
	return electrolytes[iElectrolyte];
}

const unsigned AqueousPhase::GetIndexOfElectrolyte(string nameOfElectrolyte) const
{
	for(unsigned iElectrolyte = 0; iElectrolyte < GetNumOfElectrolytes(); ++iElectrolyte)
		if(GetNameOfSpecies(GetIndexOfElectrolyte(iElectrolyte)) == nameOfElectrolyte)
			return iElectrolyte;
	
	return unsigned(-1);
}

const vector<pair<unsigned,double>> AqueousPhase::GetIonsOfElectrolyte(unsigned iElectrolyte) const
{
	return ionsOfElectrolytes[iElectrolyte];
}

const int AqueousPhase::GetElectricalChargeOfSpecies(unsigned iSpecies) const
{
	return z[iSpecies];
}

const double AqueousPhase::GetMolalityOfSpecies(unsigned iSpecies) const
{
	return 55.508 * GetMolesOfSpecies(iSpecies)/GetMolesOfSpecies(indexOfWater);
}

const double AqueousPhase::GetEffectiveIonicStrength() const
{
	double effectiveIonicStrength = 0.0;

	for(unsigned iIon = 0; iIon < GetNumOfIons(); ++iIon)
		effectiveIonicStrength += z[GetIndexOfIon(iIon)] * z[GetIndexOfIon(iIon)] * GetMolalityOfSpecies(GetIndexOfIon(iIon));
	
	return 0.5 * effectiveIonicStrength;
}

const double AqueousPhase::GetStoichiometricIonicStrength() const
{
	double stoichiometricIonicStrength = 0.0;
	
	const vector<double> mTi = GetStoichiometricMolalityOfIons();

	for(unsigned iIon = 0; iIon < GetNumOfIons(); ++iIon)
		stoichiometricIonicStrength += z[GetIndexOfIon(iIon)] * z[GetIndexOfIon(iIon)] * mTi[iIon];

	return 0.5 * stoichiometricIonicStrength;
}

const vector<double> AqueousPhase::GetStoichiometricMolalityOfIons() const
{
	// the stoichiometric molality of the ions in the aqueous phase
	vector<double> mTi(GetNumOfIons());
	
	for(unsigned iIon = 0; iIon < GetNumOfIons(); ++iIon) 
		mTi[iIon] = GetMolalityOfSpecies(GetIndexOfIon(iIon));
	
	for(unsigned iElectrolyte = 0; iElectrolyte < GetNumOfElectrolytes(); ++iElectrolyte) 
	{
		const double molalityOfElectrolyte = GetMolalityOfSpecies(GetIndexOfElectrolyte(iElectrolyte));
		
		for(unsigned iter = 0; iter < ionsOfElectrolytes[iElectrolyte].size(); ++iter)
		{
			const unsigned iIon = ionsOfElectrolytes[iElectrolyte][iter].first;
			
			const double stoichiometryOfIon = ionsOfElectrolytes[iElectrolyte][iter].second;
			
			mTi[iIon] += stoichiometryOfIon * molalityOfElectrolyte;
		}
	}
	
	return mTi;
}

const double AqueousPhase::GetTotalMolalityOfSolutes() const
{
	double aux = 0.0;
	
	for(unsigned iSpecies = 0; iSpecies < GetNumOfSpecies(); ++iSpecies)
		aux += GetMolalityOfSpecies(iSpecies);
	
	return aux - 55.508; // subtract the molality of water, which is not a solute
}

const vector<double> AqueousPhase::ActivityCoefficients(double T, double P) const
{
	vector<double> activityCoeffs(GetNumOfSpecies());
	
	ParamsHKF params;
	
	params.I = GetEffectiveIonicStrength();
	
	params.mT = GetTotalMolalityOfSolutes();
	
	params.mTi = GetStoichiometricMolalityOfIons();
	
	params.A = DebyeHuckelParameterA(T, P);
	
	params.B = DebyeHuckelParameterB(T, P);
	
	params.bNaCl = HKFParameterNaCl(T, P);
	
	params.bNapClm = HKFParameterNapClm(T, P);
	
	params.effRadii = effRadii;
	
	params.zOfIons.resize(GetNumOfIons()); 
	
	for(unsigned iIon = 0; iIon < GetNumOfIons(); ++iIon) 
		params.zOfIons[iIon] = z[GetIndexOfIon(iIon)];
	
	// the molar fraction of water in the solution (it is used to convert the activity coefficients below from mole-fraction scale to molality scale) 
	const double Xw = 55.508/(55.508 + params.mT);
	
	for(unsigned iSpecies = 0; iSpecies < GetNumOfSpecies(); ++iSpecies)
	{
		switch(activityModels[iSpecies])
		{
		case AQUEOUSACTIVITYMODEL_IDEAL:
			activityCoeffs[iSpecies] = Xw * ActivityCoefficientIdealModel();
			break;
			
		case AQUEOUSACTIVITYMODEL_HKF_ION:
			activityCoeffs[iSpecies] = Xw * ActivityCoefficientIonicModelHKF(species2ions[iSpecies], params);
			break;
			
		case AQUEOUSACTIVITYMODEL_HKF_WATER:
			activityCoeffs[iSpecies] = Xw/55.508 * ActivityWaterModelHKF(params);
			break;
			
		case AQUEOUSACTIVITYMODEL_SETSCHENOW:
			activityCoeffs[iSpecies] = Xw * ActivityCoefficientNeutralModelSetschenow(params.I, setschenowConstants[iSpecies]);
			break;
			
		default:
			activityCoeffs[iSpecies] = Xw * ActivityCoefficientNeutralModelSetschenow(params.I, setschenowConstants[iSpecies]);
			break;
		}
	}
	
	return activityCoeffs;
}

const vector<double> AqueousPhase::Activities(double T, double P) const
{
	vector<double> activities(ActivityCoefficients(T, P));
	
	const double molalityFactor = 55.508/GetMolesOfSpecies(indexOfWater);
	
	for(unsigned iSpecies = 0; iSpecies < GetNumOfSpecies(); ++iSpecies)
		activities[iSpecies] *= molalityFactor * GetMolesOfSpecies(iSpecies);
	
	return activities;
}
